
__version__ = '0.1.11'
url = 'https://gitee.com/Haskely/Python-Tools'

import sys
import socket
import random
import codecs
import argparse
import json
import os
try:
    import secrets
    def randint(n):
        return secrets.randbits(n)
except ImportError:
    def randint(n):
        return random.getrandbits(n)




class PynatError(Exception):
    """ Raised when an error occurs during network discovery. """


# Non-NAT network topologies
BLOCKED = 'Blocked'
OPEN = 'Open'
UDP_FIREWALL = 'UDP Firewall'
# NAT topologies
FULL_CONE = 'Full-cone NAT'
RESTRICTED_CONE = 'Restricted-cone NAT'
RESTRICTED_PORT = 'Restricted-port NAT'
SYMMETRIC = 'Symmetric NAT'

# Stun message types
BIND_REQUEST_MSG = b'\x00\x01'
BIND_RESPONSE_MSG = b'\x01\x01'
MAGIC_COOKIE = b'\x21\x12\xA4\x42'

# Stun attributes
MAPPED_ADDRESS = b'\x00\x01'
RESPONSE_ADDRESS = b'\x00\x02'
CHANGE_REQUEST = b'\x00\x03'
SOURCE_ADDRESS = b'\x00\x04'
CHANGED_ADDRESS = b'\x00\x05'
XOR_MAPPED_ADDRESS = b'\x00\x20'

# List of classic STUN servers
STUN_SERVERS = [('stun.ekiga.net', 3478), 
                ('stun.ideasip.com', 3478), 
                ('stun.voiparound.com', 3478),
                ('stun.voipbuster.com', 3478), 
                ('stun.voipstunt.com', 3478), 
                ('stun.voxgratia.org', 3478),
                ('stun1.l.google.com',19302),
                ('stun2.l.google.com',19302),
                ('stun3.l.google.com',19302),
                ('stun4.l.google.com',19302),
                ('23.21.150.121',3478),
                ('stun01.sipphone.com',3478),
                ('stun.ekiga.net',3478),
                ('stun.fwdnet.net',3478),
                ('stun.ideasip.com',3478),
                ('stun.iptel.org',3478),
                ('stun.rixtelecom.se',3478),
                ('stun.schlund.de',3478),
                ('stunserver.org',3478),
                ('stun.softjoys.com',3478),
                ('stun.voiparound.com',3478),
                ('stun.voipbuster.com',3478),
                ('stun.voipstunt.com',3478),
                ('stun.voxgratia.org',3478),
                ('stun.xten.com',3478),
                ]
def read_servers_file(file_path):
    with open(file_path,'r',encoding='utf-8') as f:
        SERVERS = []
        for addr_str in f.read().split('\n'):
            if addr_str:
                addr = addr_str.split(':')
                addr[1] = int(addr[1])
                SERVERS.append(tuple(addr))
    return SERVERS
def write_servers(Servers,file_path):
    Servers = set(Servers)
    with open(file_path,'w',encoding='utf-8') as f:
        f.write('\n'.join([addr[0]+":"+str(addr[1]) for addr in Servers]))

STUN_SERVERS_File_path = 'STUN_SERVERS.txt'
if not os.path.exists(STUN_SERVERS_File_path):
    write_servers(STUN_SERVERS,STUN_SERVERS_File_path)


def ORD(ch):   # compatible to python3
    return ch if type(ch) == int else ord(ch)


def long_to_bytes(n, length):  # compatible to PY2 and PY3
    # Equivalent to n.to_bytes(length,byteorder='big') in Python 3
    return bytes(bytearray((n >> i*8) & 0xff for i in range(length-1, -1, -1)))


# Send a STUN message to a server, with optional extra data
def send_stun_message(sock, addr, msg_type, trans_id=None, send_data=b''):
    if trans_id is None:
        trans_id = long_to_bytes(randint(128), 16)
    msg_len = long_to_bytes(len(send_data), 2)
    data = msg_type+msg_len+trans_id+send_data
    sock.sendto(data, addr)
    return trans_id


# Get a STUN Binding response from a server, with optional extra data
def get_stun_response(sock, addr, trans_id=None, send_data=b'', max_timeouts=6):
    timeouts = 0
    response = None
    old_timeout = sock.gettimeout()
    sock.settimeout(0.5)
    while timeouts < max_timeouts:
        try:
            trans_id = send_stun_message(sock, addr, BIND_REQUEST_MSG, trans_id, send_data)
            recv, addr = sock.recvfrom(2048)  
        except socket.timeout:
            timeouts += 1
            continue
        else:
            # Too short, not a valid message
            if len(recv) < 20:
                continue
            msg_type, recv_trans_id, attrs = recv[:2], recv[4:20], recv[20:]
            msg_len = int(codecs.encode(recv[2:4], 'hex'), 16)
            if msg_len != len(attrs):
                continue
            if msg_type != BIND_RESPONSE_MSG:
                continue
            if recv_trans_id != trans_id:
                continue
            response = {
                'ext_ip':None,
                'ext_port':None,
                'src_ip':None,
                'src_port':None,
                'change_ip':None,
                'change_port':None
            }
            i = 0
            while i < msg_len:
                attr_type, attr_length = attrs[i:i+2], int(codecs.encode(attrs[i+2:i+4], 'hex'), 16)
                attr_value = attrs[i+4:i+4+attr_length]
                i += 4 + attr_length
                if attr_length % 4 != 0:  # If not on a 32-bit boundary, add padding bytes
                    i += 4 - (attr_length % 4)
                if attr_type in [MAPPED_ADDRESS, SOURCE_ADDRESS, CHANGED_ADDRESS]:
                    family, port = ORD(attr_value[1]), int(codecs.encode(attr_value[2:4], 'hex'), 16)
                    ip = socket.inet_ntoa(attr_value[4:8])
                    if family == 0x01:  # IPv4
                        if attr_type == XOR_MAPPED_ADDRESS:
                            cookie_int = int(codecs.encode(MAGIC_COOKIE, 'hex'), 16)
                            port ^= cookie_int >> 16
                            ip = int(codecs.encode(attr_value[4:8], 'hex'), 16) ^ cookie_int
                            ip = socket.inet_ntoa(long_to_bytes(ip, 4))
                            response['xor_ip'], response['xor_port'] = ip, port
                        elif attr_type == MAPPED_ADDRESS:
                            response['ext_ip'], response['ext_port'] = ip, port
                        elif attr_type == SOURCE_ADDRESS:
                            response['src_ip'], response['src_port'] = ip, port
                        elif attr_type == CHANGED_ADDRESS:
                            response['change_ip'], response['change_port'] = ip, port
                    else:  # family == 0x02:  # IPv6
                        pass
            # Prefer, when possible, to use XORed IPs and ports
            xor_ip, xor_port = response.get('xor_ip', None), response.get('xor_port', None)
            if xor_ip is not None:
                response['ext_ip'] = xor_ip
            if xor_port is not None:
                response['ext_port'] = xor_port
            break
    sock.settimeout(old_timeout)
    return response


# Retrieve the internal working IP used to access the Internet 利用DNS服务获取本地IP
def get_internal_ip(test_addr=('114.114.114.114', 80)):  # By default, queries 114 DNS

    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.connect(test_addr)
    ip = sock.getsockname()[0]
    sock.close()
    return ip

def stun_communicate(sock, addr,typ=0):
    if addr[0] is None:
        return None
    addr = (socket.gethostbyname(addr[0]),addr[1])
    datas = [
        b'',# Get a STUN binding response from a server, without any CHANGE_REQUEST flags 与STUN服务器通讯，回复使用相同地址
        CHANGE_REQUEST + b'\x00\x04' + b'\x00\x00\x00\x02',# Get a STUN binding response from a server, asking it to change just the port from which it replies 与STUN服务器通讯，要求其回复时只改变Port
        CHANGE_REQUEST + b'\x00\x04' + b'\x00\x00\x00\x06'# Get a STUN binding response from a server, asking it to change both the IP & port from which it replies 与STUN服务器通讯，要求其回复时同时改变IP和Port
        
    ]
    if typ in (0,1,2):
        send_data = datas[typ]
    else:
        raise Exception("typ=%s不存在,typ只包含1、2、3")
    response = get_stun_response(sock, addr, send_data=send_data)
    if response is not None:
        response['send_addr'] = addr
    return response


# Find a working classic STUN server from the list
def find_stun_server(sock):

    Working_STUN_Path = "Working_STUN_SERVERS.txt"
    if os.path.exists(Working_STUN_Path):
        Working_STUN_SERVERS = read_servers_file(Working_STUN_Path)
    else:
        print("正在测试可用的STUN_Server...")
        STUN_SERVERS = read_servers_file(STUN_SERVERS_File_path)
        Working_STUN_SERVERS = []
        for stun_addr in STUN_SERVERS:
            try:
                response = get_stun_response(sock, stun_addr, max_timeouts=1)
            except socket.gaierror:  # Host not found error
                continue
            else:
                if response is not None and response['src_ip'] is not None:  # Have found a working server
                    Working_STUN_SERVERS.append(stun_addr)
        write_servers(Working_STUN_SERVERS,Working_STUN_Path)
    print("读取可用STUN Server完成,共%d个服务器!\n"%len(Working_STUN_SERVERS))
    return Working_STUN_SERVERS


def test_one_stun(sock, stun_host, stun_port, source_ip, source_port):

    Info = {
        'is_online1':None,
        'is_online2':None,
        'is_no_NAT':None,
        'is_NAT_change_ip':None,
        'is_NAT_change_port':None,
        'is_sameip_otherport_can_comein':None,
        'is_otherip_otherport_can_comein':None,
    }

    Response = {
        'response1':None,
        'response2':None,
        'sameip_otherport_response':None,
        'otherip_otherport_response':None
    }
    # 第一次连接
    stun_addr = (stun_host, stun_port)
    Response['response1'] = stun_communicate(sock, stun_addr)
    Info['is_online1'] = (Response['response1'] is not None)
    if not Info['is_online1']:
        print("该STUN服务器%s连接失败"%str(stun_addr))
        return None,Info,Response
    # 分析结果
    Info['is_no_NAT'] =  (Response['response1']['ext_ip'] == source_ip and Response['response1']['ext_port'] == source_port)
    Response['sameip_otherport_response'] = stun_communicate(sock, stun_addr,1)
    Info['is_sameip_otherport_can_comein'] =  (Response['sameip_otherport_response'] is not None)
    Response['otherip_otherport_response'] = stun_communicate(sock, stun_addr,2)
    Info['is_otherip_otherport_can_comein'] =  (Response['otherip_otherport_response'] is not None)

    # 连接第二个地址
    change_addr = (Response['response1']['change_ip'],Response['response1']['change_port'])
    Response['response2'] = stun_communicate(sock, change_addr)
    Info['is_online2'] = (Response['response2'] is not None)
    if not Info['is_online2']:
        print("该STUN服务器%s第二地址%s连接失败,这个STUN服务器有问题"%(stun_addr,change_addr))
        return None,Info,Response
    # 分析更多信息
    Info['is_NAT_change_ip'] = (Response['response1']['ext_ip'] != Response['response2']['ext_ip'])
    Info['is_NAT_change_port'] = (Response['response1']['ext_port'] != Response['response2']['ext_port'])

    if not (Info['is_online1'] and Info['is_online2']):
        topology = BLOCKED
    
    elif Info['is_no_NAT'] :
        if Info['is_otherip_otherport_can_comein']:
            topology = OPEN
        else:
            topology = UDP_FIREWALL
    else:
        if Info['is_NAT_change_ip']:
            topology = '没见过的NAT类型'
        elif Info['is_NAT_change_port']:
            topology = SYMMETRIC
        else:
            if Info['is_otherip_otherport_can_comein']:
                topology = FULL_CONE
            elif Info['is_sameip_otherport_can_comein']:
                topology = RESTRICTED_CONE
            else:
                topology = RESTRICTED_PORT
    
    def make_text(response,source_addr):
        text = ''
        if response:
            text = "发送的地址：%s\t服务器地址：%s\t自己外网地址：%s\t自己内网地址：%s"%(
                str(response['send_addr']),
                str((response['src_ip'], response['src_port'])),
                str((response['ext_ip'], response['ext_port'])),
                source_addr
                )
        else:
            text = '没有回应'
        return text
    source_addr= (source_ip,source_port)
    res_text = ''
    for res_name,res in Response.items():
        res_text += make_text(res,source_addr) + '\t' + res_name +'\n'
    print(res_text + 'Topology:%s,STUN地址:%s'%(topology,stun_addr)+'\n')

    return topology,Info,Response

import traceback
# Get the network topology, external IP, and external port
def get_ip_info(source_ip='0.0.0.0', source_port=54321, stun_host=None, stun_port=3478,
                sock=None):
    """ Get information about the network topology, external IP, and external port.

    Args:
        source_ip (str, optional): If not '0.0.0.0', the internal IP address to bind to. Defaults to '0.0.0.0'.
        source_port (int, optional): Source port to bind to. Defaults to 54320.
        stun_host (str, optional): Address of the STUN host to use. Defaults to None, in which case one is selected.
        stun_port (int, optional): Port of the STUN host to query. Defaults to 3478.
        include_internal (bool, optional): Whether to include internal IP address information. Defaults to False.
        sock (socket.socket, optional): Bound socket to connect with. If not provided, one will be created.

    Returns:
        tuple: (topology, external_ip, external_port). Topology & external_ip are strings, external_port is int.
            If `include_internal` is True, returns (topology, external_ip, external_port, internal_ip) instead.
    """
    # If no socket is passed in, create one and close it when done
    ephemeral_sock = sock is None
    if sock is None:
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        sock.bind((source_ip, source_port))
    try:
        temp_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        temp_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        temp_sock.connect(('114.114.114.114',80))
    except:
        traceback.print_exc()
        print("网络不连通")
        return BLOCKED
    finally:
        temp_sock.close()
    # Determine the actual local, or source IP
    if source_ip == '0.0.0.0':
        source_ip = get_internal_ip()
    # Find a stun host if none was selected
    Topologys = []
    Infos = []
    Responses = []
    if stun_host is None:
        stun_addrs = find_stun_server(sock)
        # If None was found, assume the network is blocked
        if len(stun_addrs) == 0:
            print("没有可用的STUN服务器")
            topology = BLOCKED
        else:
            Tested_stun_addrs = set()
            for stun_addr in stun_addrs:
                pure_addr = (socket.gethostbyname(stun_addr[0]),stun_addr[1])
                if pure_addr not in Tested_stun_addrs:
                    Tested_stun_addrs.add(pure_addr)
                    topology,Info,Response = test_one_stun(sock,stun_addr[0],stun_addr[1],source_ip,source_port)
                    Topologys.append(topology)
                    Infos.append(Info)
                    Responses.append(Response)
                else:
                    print("[STUN地址:%s对s应的ip%s已经被测试过了，跳过]\n"%(str(stun_addr),str(pure_addr)))
    # If a stun hose was specified, set stun_addr
    else:
        topology,Info,Response = test_one_stun(sock,stun_host,stun_port,source_ip,source_port)
        Topologys.append(topology)
        Infos.append(Info)
        Responses.append(Response)

    if ephemeral_sock:
        sock.close()

    def summary(Topologys,Infos,Responses):
        print("===Summary===")
        titles = ['*****发送的地址*****','*****回复的地址*****','*****自己外网地址*****','*****自己内网地址*****']
        datas = []
        for Res in Responses:
            for res in Res.values():
                if res:
                    datas.append((res['send_addr'],(res['src_ip'], res['src_port']),(res['ext_ip'], res['ext_port']),(source_ip,source_port)))
        datas = list(set(datas))
        datas.sort()
        print('\t'.join(titles))
        for data in datas:
            print('\t'.join([str(d) for d in data]))
    summary(Topologys,Infos,Responses)
    return topology




def main():
    try:
        parser = argparse.ArgumentParser(prog='pynat', description=__doc__)
        parser.add_argument('--source_ip', help='The source IPv4 address to bind to.', type=str, default='0.0.0.0')
        parser.add_argument('--source-port', help='The source port to bind to.', type=int, default=54321)
        parser.add_argument('--stun-host', help='The STUN host to use for queries.', type=str)
        parser.add_argument('--stun-port', help='The port of the STUN host to use for queries.', type=int, default=3478)
        args = parser.parse_args()
        source_ip, source_port, stun_host, stun_port = args.source_ip, args.source_port, args.stun_host, args.stun_port
        topology = get_ip_info(source_ip, source_port, stun_host, stun_port)
        print('Network type:', topology)
    except KeyboardInterrupt:
        sys.exit()
    else:
        sys.exit(0)


if __name__ == '__main__':
    main()
